import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation

# 定义两个复数
r1 = 2
theta1 = np.pi / 4
r2 = 1
theta2 = np.pi / 6

# 计算复数的实部和虚部
x1 = r1 * np.cos(theta1)
y1 = r1 * np.sin(theta1)
x2 = r2 * np.cos(theta2)
y2 = r2 * np.sin(theta2)

# 计算乘积的模和辐角
r_product = r1 * r2
theta_product = theta1 + theta2
x_product = r_product * np.cos(theta_product)
y_product = r_product * np.sin(theta_product)

# 创建 3D 图形
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

# 设置坐标轴范围
ax.set_xlim([-3, 3])
ax.set_ylim([-3, 3])
ax.set_zlim([0, 3])
ax.set_xlabel('Real Axis')
ax.set_ylabel('Imaginary Axis')
ax.set_zlabel('Magnitude')

# 初始化向量
vector1, = ax.plot([0, x1], [0, y1], [0, r1], color='r', label='$z_1$')
vector2, = ax.plot([0, x2], [0, y2], [0, r2], color='g', label='$z_2$')
product_vector, = ax.plot([0, 0], [0, 0], [0, 0], color='b', label='$z_1z_2$')

# 定义动画更新函数
def update(frame):
    # 逐渐展示乘积向量
    alpha = frame / 100
    current_x = alpha * x_product
    current_y = alpha * y_product
    current_r = alpha * r_product
    product_vector.set_data([0, current_x], [0, current_y])
    product_vector.set_3d_properties([0, current_r])
    return product_vector,

# 创建动画
ani = FuncAnimation(fig, update, frames=100, interval=50, blit=True)

# 显示图例
ax.legend()

# 显示动画
plt.show()
